package com.sunday.common.mq.rocket.brave;

import brave.Span;
import brave.Tracer;
import brave.Tracing;
import brave.internal.Nullable;
import brave.messaging.MessagingRequest;
import brave.messaging.MessagingTracing;
import brave.propagation.Propagation;
import brave.propagation.TraceContext;
import brave.propagation.TraceContext.Extractor;
import brave.propagation.TraceContext.Injector;
import brave.propagation.TraceContextOrSamplingFlags;
import brave.sampler.SamplerFunction;
import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.client.consumer.DefaultLitePullConsumer;
import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
import org.apache.rocketmq.client.impl.consumer.DefaultLitePullConsumerImpl;
import org.apache.rocketmq.client.producer.DefaultMQProducer;
import org.apache.rocketmq.client.trace.TraceConstants;
import org.apache.rocketmq.common.message.Message;

import java.lang.reflect.Field;
import java.util.Map;

/**
 * spring rocket tracing 兼容
 *
 * @author zsy
 * @see brave.spring.rabbit.SpringRabbitTracing 参考
 * @since 2023/7/31
 */
@Slf4j
public final class RocketTracing {

    static final String ROCKETMQ_TOPIC = "rocketmq.topic";

    public static RocketTracing create(Tracing tracing) {
        return newBuilder(tracing).build();
    }

    public static RocketTracing create(MessagingTracing messagingTracing) {
        return newBuilder(messagingTracing).build();
    }

    public static Builder newBuilder(Tracing tracing) {
        return newBuilder(MessagingTracing.create(tracing));
    }

    /**
     * @since 5.9
     */
    public static Builder newBuilder(MessagingTracing messagingTracing) {
        return new Builder(messagingTracing);
    }

    public static final class Builder {
        final MessagingTracing messagingTracing;
        String remoteServiceName = TraceConstants.ROCKETMQ_SERVICE;

        Builder(MessagingTracing messagingTracing) {
            if (messagingTracing == null) throw new NullPointerException("messagingTracing == null");
            this.messagingTracing = messagingTracing;
        }

        public Builder remoteServiceName(String remoteServiceName) {
            this.remoteServiceName = remoteServiceName;
            return this;
        }

        public RocketTracing build() {
            return new RocketTracing(this);
        }
    }

    final Tracing tracing;
    final Tracer tracer;
    final Extractor<MessageProducerRequest> producerExtractor;
    final Extractor<MessageConsumerRequest> consumerExtractor;
    final Injector<MessageProducerRequest> producerInjector;
    final Injector<MessageConsumerRequest> consumerInjector;
    final String[] traceIdHeaders;
    final SamplerFunction<MessagingRequest> producerSampler, consumerSampler;
    final String remoteServiceName;
    @Nullable
    final Field defaultLitePullConsumerImplField;

    RocketTracing(Builder builder) {
        this.tracing = builder.messagingTracing.tracing();
        this.tracer = tracing.tracer();
        MessagingTracing messagingTracing = builder.messagingTracing;
        Propagation<String> propagation = messagingTracing.propagation();
        this.producerExtractor = propagation.extractor(MessageProducerRequest.GETTER);
        this.consumerExtractor = propagation.extractor(MessageConsumerRequest.GETTER);
        this.producerInjector = propagation.injector(MessageProducerRequest.SETTER);
        this.consumerInjector = propagation.injector(MessageConsumerRequest.SETTER);
        this.producerSampler = messagingTracing.producerSampler();
        this.consumerSampler = messagingTracing.consumerSampler();
        this.remoteServiceName = builder.remoteServiceName;
        this.traceIdHeaders = propagation.keys().toArray(new String[0]);
        defaultLitePullConsumerImplField =
                getField(DefaultLitePullConsumer.class, "defaultLitePullConsumerImpl");
    }

    @Nullable
    static Field getField(Class<?> clazz, String name) {
        Field result = null;
        try {
            result = clazz.getDeclaredField(name);
            result.setAccessible(true);
        } catch (NoSuchFieldException e) {
        }
        return result;
    }

    protected SendMessageBraveTracingHookImpl newBraveSendMessageHook() {
        return new SendMessageBraveTracingHookImpl(this);
    }

    public DefaultMQProducer decorateDefaultMQProducer(DefaultMQProducer producer) {
        producer.getDefaultMQProducerImpl()
                .registerSendMessageHook(newBraveSendMessageHook());
        return producer;
    }

    protected ConsumeMessageBraveTracingHookImpl newBraveConsumeMessageHook() {
        return new ConsumeMessageBraveTracingHookImpl(this);
    }

    public DefaultMQPushConsumer decorateDefaultMQPushConsumer(DefaultMQPushConsumer consumer) {
        consumer.getDefaultMQPushConsumerImpl()
                .registerConsumeMessageHook(newBraveConsumeMessageHook());
        return consumer;
    }

    protected PullConsumeMessageBraveTracingHookImpl newBravePullConsumeMessageHook() {
        return new PullConsumeMessageBraveTracingHookImpl(this);
    }

    public DefaultLitePullConsumer decorateDefaultLitePullConsumer(DefaultLitePullConsumer pullConsumer) {
        if (defaultLitePullConsumerImplField == null) return pullConsumer;
        try {
            // don't use "field.get(obj) instanceof X" as the field could be null
            if (DefaultLitePullConsumerImpl.class.isAssignableFrom(defaultLitePullConsumerImplField.getType())) {
                DefaultLitePullConsumerImpl defaultLitePullConsumerImpl =
                        (DefaultLitePullConsumerImpl) defaultLitePullConsumerImplField.get(pullConsumer);
                defaultLitePullConsumerImpl.registerConsumeMessageHook(newBravePullConsumeMessageHook());
            }
        } catch (Exception e) {
        }
        return pullConsumer;
    }

    <R> TraceContextOrSamplingFlags extractAndClearTraceIdHeaders(
            Extractor<R> extractor, R request, Message message
    ) {
        // 返回跟踪上下文或从请求解析的采样标志。如果什么都没有可解析的，采样标志将设置为{@link SamplingFlags#EMPTY}。
        TraceContextOrSamplingFlags extracted = extractor.extract(request);
        /**
         * 当不存在{@linkplain TraceContext#traceIdString() trace ID}时返回非{@code null}
         * {@link Extractor#extract(Object) extracted}从一个请求。
         * 例如，给定标题“b3: 1”，{@link B3Propagation}提取{@link #SAMPLED}。
         * 当{@link #context()}和{@link #traceIdContext()}不是{@code null}
         */
        // 清除标头中存在的所有传播键
        if (extracted.samplingFlags() == null) { // 然后提取跟踪id
            Map<String, String> properties = message.getProperties();
            if (properties != null) clearTraceIdHeaders(properties);
        }
        return extracted;
    }

    /**
     * 我们不能因为可能注入单个B3就跳过清除头，而得到陈旧的B3多个，反之亦然。
     */
    void clearTraceIdHeaders(Map<String, String> headers) {
        for (String traceIDHeader : traceIdHeaders) headers.remove(traceIDHeader);
    }

    /**
     * 创建表示此请求的潜在noop远程跨度
     */
    Span nextMessagingSpan(
            SamplerFunction<MessagingRequest> sampler,
            MessagingRequest request,
            TraceContextOrSamplingFlags extracted
    ) {
        /**
         * {@link TraceContextOrSamplingFlags#sampled(boolean)}
         * 根据具体的实现和上下文，TraceContextOrSamplingFlags.sampled()
         * 方法可能返回一个布尔值，指示给定的跟踪上下文或采样标志是否被采样。
         * 如果方法返回 true，则表示跟踪上下文或采样标志已经被采样；
         * 如果返回 false，则表示未被采样。
         */
        Boolean sampled = extracted.sampled();
        // 只有在消息传递采样器做出决定时才重新创建上下文
        if (sampled == null && (sampled = sampler.trySample(request)) != null) {
            extracted = extracted.sampled(sampled.booleanValue());
        }
        return tracer.nextSpan(extracted);
    }

}
